#pragma once
#include <iostream>
#include <fstream>
#include <string>
#include <cstdlib>
#include <map>
#include "codec.cpp"

using std::cout;
using std::endl;
using std::fstream;
using std::map;
using std::ios;
using std::make_pair;
using std::vector;

vector<pair<char, int> > getWeight_en(string &);
void encode_file(string &, string &, vector<inf> &, vector<pair<char, int> > &);
void writeBit(fstream &, bool, bool flush = false);

// 编码主程序
void encode_main(string infname, string oufname)
{
	vector<pair<char, int> > weight;
	vector<pair<char, int> >::iterator weight_p;
	vector<inf> dict;
	vector<inf>::iterator dict_p;
	HuffmanTree h_tree;

	weight = getWeight_en(infname);									//统计文本文件中各个字母出现的次数
	
	heapSort(weight);												//将所有字符按权值排序

	h_tree = BuildHuffmanTree(weight);								//生成哈夫曼树
	h_tree.CreateDict(dict);										//生成编码字典

	for(dict_p = dict.begin(); dict_p != dict.end(); dict_p++)		//在屏幕上输出编码
	{
		string str;
		char ch = dict_p->ch;
		if(ch == '\n')
			str = "(LF)";
		else if(ch == '\t')
			str = "(TAB)";
		else if(ch == ' ')
			str = "(SPACE)";
		else
			str = ch;
		cout << str << "\t" << dict_p->weight << "\t" << dict_p->code << endl;
	}

	encode_file(infname, oufname, dict, weight);					//将原文件编码后写入输出文件

	return;
}

//统计输入的文本文件中各个字符出现的次数
vector<pair<char, int> > getWeight_en(string &infname)
{
	vector<pair<char, int> > weight(128);							//做一个“桶”
	vector<pair<char, int> >::iterator p, t;
	char buffer;

	fstream fin(infname.c_str(), ios::in);							//打开文本文件
	if(fin.fail())
	{
		cout << "Cannot open the input file." << endl;
		exit(0);
	}

	for(int i = 0; i < 128; i++)
		weight[i] = make_pair((char)i, 0);

	while((buffer = fin.get()) != EOF)
		weight[(int)buffer].second++;								//遍历文件，统计文件中各个字符出现的次数

	for(p = weight.begin(); p != weight.end();)
		if(((*p).second) == 0)
		 	p = weight.erase(p);									//将权值列表中出现次数为0的字符去掉
		else
			p++;

	fin.close();
	return weight;
}

//按位写文件函数
void writeBit(fstream &fp, bool bit, bool flush)
{
	static char buffer = 0;
	static int offset = 0;											//offset表示当前buffer中已写入的位数
	if(flush && offset != 0)										//当编码完成后，将buffer中剩余的二进制位写入文件
	{
		while(offset != 8)
		{
			buffer <<= 1;
			offset++;
		}
		fp.write(&buffer, 1);
		// fp.flush();
		return;
	}

	buffer <<= 1;
	offset++;
	buffer |= (char)bit;											//将新接收到的位写入buffer的低位
	if(offset == 8)													//若buffer已满，则将buffer写入文件
	{
		fp.write(&buffer, 1);
		buffer = 0;
		offset = 0;
	}
}

//文件编码函数
void encode_file(string &infname, string &oufname, vector<inf> &dict, vector<pair<char, int> > &weight)
{
	fstream fin(infname.c_str(), ios::in);
	fstream fout(oufname.c_str(), ios::out|ios::binary);
	map<char, string> ch_dict;
	vector<inf>::iterator d_p;
	vector<pair<char, int> >::iterator w_p;
	char temp;
	string coding;
	int size = dict.size(), code_size = 0;

	if(fin.fail())
	{
		cout << "Cannot open the input file." << endl;
		exit(0);
	}
	if(fout.fail())
	{
		cout << "Cannot open the output file." << endl;
		exit(0);
	}

	fout.write((const char *)&size, sizeof(int));					//为码长留出空间，方便最后统计码长后写入文件头
	fout.write((const char *)&size, sizeof(int));

	for(d_p = dict.begin(); d_p != dict.end(); d_p++)
		ch_dict.insert(make_pair(d_p->ch, d_p->code));				//构造映射字典

	for(w_p = weight.begin(); w_p != weight.end(); w_p++)			//将权值字典写入编码文件
	{
		fout.write((const char *)&(w_p->first), sizeof(char));		//写入字符
		fout.write((const char *)&(w_p->second), sizeof(int));		//写入字符出现的次数
	}

	while((temp = fin.get()) != EOF)								//遍历文本文件
	{
		coding = ch_dict.find(temp)->second;						//找到所读字符对应的哈夫曼编码
		code_size += coding.length()+1;
		for(unsigned int i = 0; i < coding.length(); i++)
			writeBit(fout, (coding[i] == '1'));						//将哈夫曼编码按位写入编码文件
	}
	writeBit(fout, true, true);										//将缓冲区中剩余的数据写入编码文件

	fout.seekp(0L, ios::beg);
	fout.write((const char *)&code_size, sizeof(int));				//将编码总长度写入文件头

	cout << "code size: " << code_size << endl;
	cout << "Encoding finished" << endl;
	fin.close();
	fout.close();													//编码文件格式：总码长，字符种类，n个字符及权值，已编码的文本
	return;
}